import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split


def load_iris_binary():
    """
    XGBoost 二分类 预测是否为“virginica”
    特征数：4个（花萼长度、花萼宽度、花瓣长度、花瓣宽度）
    样本数：150个
    """
    # 加载数据
    iris = load_iris()
    X, y = iris.data, iris.target

    # 将问题转换为二分类：是否为virginica (类别2)
    y_binary = (y == 2).astype(int)

    # 展示数据类型
    print("X的数据类型:", type(X))
    print("y的数据类型:", type(y))
    print("X的形状:", X.shape)
    print("y的形状:", y.shape)
    print(f"X的前5行:\n{X[:5]}")
    print("非virginica (0):", np.sum(y_binary == 0))
    print("virginica (1):", np.sum(y_binary == 1))

    return X, y_binary


def load_iris_multiclass():
    """
    XGBoost 多分类示例 - 使用鸢尾花数据集
    鸢尾花数据集包含3个类别：0-setosa, 1-versicolor, 2-virginica
    特征数：4个（花萼长度、花萼宽度、花瓣长度、花瓣宽度）
    样本数：150个
    """

    # 加载数据并划分
    X, y = load_iris(return_X_y=True)

    print(f"数据形状: X{X.shape}, y{y.shape}")
    print(f"类别分布: {np.bincount(y)}")
    print(f"X的前5行:\n{X[:5]}")

    return X, y


if __name__ == "__main__":
    load_iris_multiclass()
